package karpiel.boguchwal.wojciech.przemek

import karpiel.boguchwal.wojciech.przemek.Ast._
import karpiel.boguchwal.wojciech.przemek.RealCache.InnerCache

import scala.collection.mutable


trait Cache {
  def infer(ast: Ast, orElse: () => Ast)(implicit ctx: Context): Ast

  def normalize(ast: Ast, orElse: () => Ast)(implicit ctx: Context): Ast
}


case class RealCache(cache: mutable.WeakHashMap[Context, InnerCache] = new mutable.WeakHashMap[Context, InnerCache]()) extends Cache {

  import RealCache._

  def infer(expr: Ast, orElse: () => Ast)(implicit ctx: Context): Ast =
    work(expr, ctx, orElse, PickInfer)

  def normalize(expr: Ast, orElse: () => Ast)(implicit ctx: Context): Ast =
    work(expr, ctx, orElse, PickNorm)

  private def work(expr: Ast, ctx: Context, orElse: () => Ast, picker: CachePicker): Ast = {
    val squashed = squash(expr, ctx)
    val innerCache = cache.getOrElseUpdate(squashed.ctx, InnerCache())
    val map = picker.pick(innerCache)
    val squashedResult = map.getOrElseUpdate(squashed.expr, suquashExpr(orElse(), squashed.oldToNew))
    restoreExpr(squashedResult, squashed.oldToNew)

  }

}


object DummyCache extends Cache {
  override def infer(ast: Ast, orElse: () => Ast)(implicit ctx: Context): Ast = orElse()

  override def normalize(ast: Ast, orElse: () => Ast)(implicit ctx: Context): Ast = orElse()
}


object RealCache {

  sealed trait CachePicker {
    def pick(c: InnerCache): mutable.WeakHashMap[Ast, Ast]
  }

  object PickInfer extends CachePicker {
    override def pick(c: InnerCache): mutable.WeakHashMap[Ast, Ast] = c.inferCache
  }

  object PickNorm extends CachePicker {
    override def pick(c: InnerCache): mutable.WeakHashMap[Ast, Ast] = c.normalCache
  }
  case class InnerCache(inferCache: mutable.WeakHashMap[Ast, Ast] = new mutable.WeakHashMap[Ast, Ast](),
                        normalCache: mutable.WeakHashMap[Ast, Ast] = new mutable.WeakHashMap[Ast, Ast]())

  case class Squashed(expr: Ast, ctx: Context, oldToNew: Map[Int, Int])


  def suquashExpr(expr: Ast, oldToNew: Map[Int, Int], nestedness: Int = 0): Ast = {
    expr match {
      case Bruijn(i, name) if i >= nestedness => Bruijn(nestedness + oldToNew(i - nestedness), name)
      case PiValue(typeT, value) => PiValue(suquashExpr(typeT, oldToNew, nestedness), suquashExpr(value, oldToNew, nestedness + 1))
      case PiType(typeT, value) => PiType(suquashExpr(typeT, oldToNew, nestedness), suquashExpr(value, oldToNew, nestedness + 1))
      case anyElse => Utils.applyToSubTrees(anyElse, a => suquashExpr(a, oldToNew, nestedness))
    }
  }


  def squash(expr: Ast, ctx: Context): Squashed = {

    val exprBruijns = Utils.usesBruins(expr)

    def totalUsedBruijnsQ(ctx: Context, b: Set[Int]): Set[Int] = {
      val newSet = b.flatMap(idx => {
        val element = ctx.ctx(idx)
        Utils.usesBruins(element.typeT) ++
          element.value.map(v => Utils.usesBruins(v)).getOrElse(Set())
      }
      )

      if (newSet.subsetOf(b)) {
        b
      } else {
        totalUsedBruijnsQ(ctx, newSet ++ b)
      }

    }

    val oldToNew = totalUsedBruijnsQ(ctx, exprBruijns.toSet).toSeq.sorted.zipWithIndex.toMap
    val trimmedCtx = Context(oldToNew.keys.toSeq /*toSeq musi KONIECZNIE tu być bo inaczej się robi z tego zbiór*/ .map(idx => ctx.ctx(idx)).toList)
    val squashedCtx = Context(trimmedCtx.ctx.map {
      case ContextElement(typeT, value, name) =>
        ContextElement(suquashExpr(typeT, oldToNew),
          value.map(v => suquashExpr(v, oldToNew)),
          name)
    })

    Squashed(suquashExpr(expr, oldToNew), squashedCtx, oldToNew)
  }

  def restoreExpr(expr: Ast, oldToNew: Map[Int, Int]): Ast = {
    val newToOld = oldToNew.map { case (a, b) => b -> a }
    suquashExpr(expr, newToOld)
  }
}

/*
EmptyType
BoolFalse
BoolTrue
BoolType
Universe
UnitType
UnitValue
 */